#!/usr/bin/python 

import sys
import os
import simplejson as json

__author__="qingkaishi"
__date__ ="$Apr 21, 2014 9:56:47 AM$"

class SysFlags:
    def __init__(self):
        pass
    
    @staticmethod
    def parse(map):
        sf = SysFlags()
        sf.flags = map # TODO
        return sf
    
class ChildMasterResource:
    def __init__(self):
        pass
    
    @staticmethod
    def parse(map):
        cmr = ChildMasterResource()
        cmr.cpu = map["CPU"]
        cmr.mem = map["Memory"]
        cmr.sinstance = map["SInstance"]
        
        return cmr

class Description:
    def __init__(self):
        pass
    
    def dot_job_dag(self, name):
        fileHandler = open(name+'.dot', 'w')
        fileHandler.write('digraph '+name+'{\n')
        
        idxmap = {}
        idx = 1
        for t in self.tasks:
            fileHandler.write('a'+str(idx)+'[label='+str(t.name)+'];\n')
            idxmap[t] = idx
            idx = idx +1 
            
        for p in self.pipes:
            src = p.src
            dst = p.dst
            
            if(src.type()=='Task'):
                sidx = idxmap[src]
            else:
                sidx = idx
                idx = idx + 1
                
            if(dst.type() == 'Task'):
                didx = idxmap[dst]
            else:
                didx = idx
                idx= idx + 1
            
            fileHandler.write('a'+str(sidx)+'->a'+str(didx)+';\n')
        
        fileHandler.write('}\n')
        fileHandler.close()
    
    @staticmethod
    def parse(map, statusobj):
        desc = Description()
        
        desc.handle_broken_pipe = bool(map["HandleBrokenPipe"])
        desc.import_snapshort = bool(map["ImportSnapshot"])
        desc.nuwa_capability = map["NuwaCapability"]
        desc.package_capability = map["PackageCapability"]
        desc.pangu_capability = map["PanguCapability"]
        desc.tasks = []
        desc.pipes = []
        
        nameTaskMap = {}
        rawTasks = map["Tasks"]
        for name, rt in rawTasks.items():
            t = ApsaraTask.parse(name, rt, statusobj)
            desc.tasks.append(t)
            nameTaskMap[name] = t
        
        rawPipes = map["Pipes"]
        # connect tasks using pipes, and nameTaskMap
        for item in rawPipes:
            p = ApsaraPipe.parse(item, nameTaskMap)
            desc.pipes.append(p)
        
        return desc

class ApsaraJob:
    def __init__(self):
        pass
    
    def dot_job_dag(self):
        self.description.dot_job_dag(self.name)
    
    @staticmethod
    def parse(jsonfile, statusfile):
        try:
            jobobj = json.load(open(jsonfile))
        except Exception, e:
            print e.message + ' '+ jsonfile
            jobobj = json.load(open(jsonfile).read().decode("UTF-8", "replace").encode("UTF-8"))
            
        try:
            statusobj = json.load(open(statusfile))
        except:
            statusobj = json.load(open(statusfile).read().decode("UTF-8", "replace").encode("UTF-8"))
            
        apsjob = ApsaraJob()
        apsjob.account = jobobj["Account"]
        apsjob.name = jobobj["Name"]
        apsjob.priority = jobobj["Priority"]
        apsjob.type = jobobj["Type"]
        apsjob.user = jobobj["User"]
        
        apsjob.sys_flags = SysFlags.parse(jobobj["SysFlags"])
        apsjob.child_master_resource = ChildMasterResource.parse(jobobj["ChildMasterResource"])
        apsjob.description = Description.parse(jobobj["Description"], statusobj)
        
        return apsjob
    
class DAGNode:
    def __init__(self):
        self.outgoing_edges = []
        self.ingoing_edges = []
        self.outgoing_nodes = []
        self.ingoing_nodes = []
    
    def dsts(self):
        return outgoing_nodes
    
    def srcs(self):
        return ingoing_nodes
    
    def froms(self):
        return ingoing_edges
    
    def tos(self):
        return outgoing_edges
    
    def type(self):
        pass
    
class ApsaraTask(DAGNode):
    def __init__(self):
        DAGNode.__init__(self)
        self.instance = []
        
    def type(self):
        return 'Task'
    
    @staticmethod
    def parse(name, raw, statusobj):
        at = ApsaraTask()
        
        at.name = name
        at.executable_uri = raw["ExecutableUri"]
        at.instance_count = raw["InstanceCount"]
        at.instance_running_mode = raw["InstanceRunningMode"]
        at.package_capability = raw["PackageCapability"]
        at.parameters = raw["Parameters"] # TODO
        at.resource = raw["Resource"] # TODO
        at.retry_times = raw["RetryTimes"]
        at.scheduler_type = raw["SchedulerType"]
        at.tasker_worker_class_name=raw["TaskWorkerClassName"]
        at.instances = [] # TODO
        
        return at
    
    
class ApsaraFile(DAGNode):
    def __init__(self, files):
        DAGNode.__init__(self)
        self.files = files
        self.name = "files"
        
    def type(self):
        return 'File'

class ApsaraPipe:
    def __init__(self):
        pass
    
    @staticmethod
    def parse(item, nameTaskMap):
        ap = ApsaraPipe()
        ap.mode = item['Mode']
        if item.has_key('CarrierMode'):
            ap.carrier_mode = item['CarrierMode']
        
        for srctype, srcs in item["Source"].items():
            if srctype == "AccessPoint":
                ap.src = nameTaskMap[srcs.split(":")[0]]
            elif srctype == "FilePattern":
                ap.src = ApsaraFile(srcs)
            else:
                print "Unhandled src type:" + srctype
            
        for dsttype, dsts in item["Destination"].items():
            if dsttype == "AccessPoint":
                ap.dst = nameTaskMap[dsts.split(":")[0]]
            elif dsttype == "FilePattern":
                ap.dst = ApsaraFile(dsts)
            else:
                print "Unhandled src type:" + srctype
                
        ap.src.outgoing_edges.append(ap)
        ap.dst.ingoing_edges.append(ap)
        
        ap.src.outgoing_nodes.append(ap.dst)
        ap.dst.ingoing_edges.append(ap.src)
        
        return ap

class ApsaraInstance:
    def __init__(self):
        pass


if __name__ == '__main__':
    job = ApsaraJob.parse(sys.argv[1], sys.argv[2]) # json, jobstatus
    job.dot_job_dag()